import argparse
#from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from dataset.PiFu import PiFu
import socket
from datetime import datetime
import os
from model.BaseNet_V7 import CPFNet, BaseNet_version
from model.unet import UNet
import torch
from torch.utils.tensorboard import SummaryWriter
import tqdm
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from  PIL import Image
import logging
import utils.utils as u
import utils.loss as LS
from utils.config import DefaultConfig
import torch.backends.cudnn as cudnn



def new_test(args, epoch, model, dataloader):
	# 函数已修改，适应皮肤分割的单通道输出
	# print()
	# print('Start Validation!')
	# print("dataloader_val size:", len(dataloader))
	dice = precsion = jaccard = 0
	n_val = len(dataloader)  # the number of batch
	with torch.no_grad():
		model.eval()
		tbar = tqdm.tqdm(total=len(dataloader) * args.batch_size_val)
		tbar.set_description('val %d' % epoch)
		for i, (data, labels, boundary) in enumerate(dataloader):
			# tbar.update()
			if torch.cuda.is_available() and args.use_gpu:
				data = data.cuda()
				labels = labels.cuda()
			# print("data size:", data.size())
			# print("lables size:", labels.size())
			mask_pred, mask_d1, _, _, _ = model(data)
			dice_b, precsion_b, jaccard_b = u.eval_seg(mask_pred, labels)
			dice = dice + dice_b
			precsion = precsion + precsion_b
			jaccard += jaccard_b

			tbar.update(args.batch_size_val)
			tbar.set_postfix(dice='%.6f' % (dice / (i + 1)))  # 平均batch的 dice
		tbar.close()
		print()
	model.train()
	return dice / n_val, precsion / n_val, jaccard / n_val


def main(mode='train', args=None):
	# create dataset and dataloader
	dataset_path = os.path.join(args.data, args.dataset)
	print("dataset_path :", dataset_path)
	dataset_train = PiFu(dataset_path, scale=(args.crop_height, args.crop_width), k_fold_test=args.k_fold, mode='train')
	print("训练集数量大小：", len(dataset_train))
	
	dataloader_train = DataLoader(
		dataset_train,
		batch_size=args.batch_size,
		shuffle=True,
		num_workers=args.num_workers,
		pin_memory=True,
		drop_last=True
	)
	print("训练集一个epoch的batch数量：", len(dataloader_train))
	
	dataset_val = PiFu(dataset_path, scale=(args.crop_height, args.crop_width), k_fold_test=args.k_fold, mode='val')
	dataloader_val = DataLoader(
		dataset_val,
		# this has to be 1
		# batch_size=len(args.cuda.split(',')),
		batch_size=args.batch_size_val,
		shuffle=False,
		num_workers=args.num_workers,
		pin_memory=True,
		drop_last=True
	)
	print("验证集数量大小：", len(dataset_val))
	print("验证集一个epoch的batch数量：", len(dataloader_val))
	
	# exit(0)
	dataset_test = PiFu(dataset_path, scale=(args.crop_height, args.crop_width), k_fold_test=args.k_fold, mode='test')
	dataloader_test = DataLoader(
		dataset_test,
		# this has to be 1
		# batch_size=len(args.cuda.split(',')),
		batch_size=args.batch_size_test,
		shuffle=False,
		num_workers=args.num_workers,
		pin_memory=True,
		drop_last=True
	)
	# build model
	os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
	
	# load model
	model_all = {'BaseNet': CPFNet(out_planes=args.num_classes), 'UNet': UNet(in_channels=1, n_classes=args.num_classes)}
	model = model_all[args.net_work]
	cudnn.benchmark = True
	# model._initialize_weights()
	# model.cuda()
	print("torch.cuda.is_available() :", torch.cuda.is_available(), " Use :", args.cuda)
	if torch.cuda.is_available() and args.use_gpu:
		model = torch.nn.DataParallel(model).cuda()
	
	# load pretrained model if exists
	print("=> loading pretrained model '{}'".format(args.pretrained_model_path))
	checkpoint = torch.load(args.pretrained_model_path)
	checkpoint = checkpoint['state_dict']
	model.load_state_dict(checkpoint)
	print('Done!')
	
	epoch = 0
	dice, precsion, jaccard = new_test(args, epoch, model, dataloader_test)
	# dice jacd accu epoch
	if dice > args.b_dice[0]:
		args.b_dice = [dice, jaccard, precsion, epoch]
	print("          Dice                  Ja                Acc")
	print("Best Dice", args.b_dice)


if __name__ == '__main__':
	seed = 1234
	torch.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)
	args = DefaultConfig()
	logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
	logging.info(f'''Starting training:
         NetWork:         {args.net_work}
         Epochs:          {args.num_epochs}
         Batch size:      {args.batch_size}
         Batch size_val:  {args.batch_size_val}
         Learning rate 0: {args.lr[0]}
         Learning rate 1: {args.lr[1]}
         Val_k_fold:      {args.k_fold}
         Checkpoints:     {args.save_model_path}
         Device cuda:     {args.cuda}
         ''')
	
	modes = args.mode
	args.BaseNet_version = BaseNet_version
	
	if modes == 'train':
		main(mode='train', args=args)
	elif modes == 'test':
		main(mode='test', args=args)
	elif modes == 'train_test':
		main(mode='train_test', args=args)
